# -*- coding:UTF-8 -*-

import matplotlib
matplotlib.use('TkAgg')

import matplotlib.pyplot as plt
import networkx as nx

def get_edgelist_by_weight_equal(G, weight ):
    return [(u, v) for (u, v, d) in G.edges(data=True) if d['weight'] == weight]

def get_edgelist_by_weight_greater(G, weight ):
    return [(u, v) for (u, v, d) in G.edges(data=True) if d['weight'] > weight]

def get_edgelist_by_weight_greaterequal(G, weight ):
    return [(u, v) for (u, v, d) in G.edges(data=True) if d['weight'] >= weight]

def get_edgelist_by_weight_less(G, weight ):
    return [(u, v) for (u, v, d) in G.edges(data=True) if d['weight'] < weight]

def get_edgelist_by_weight_lessequal(G, weight ):
    return [(u, v) for (u, v, d) in G.edges(data=True) if d['weight'] <= weight]

def nx_get_pos(G):
    return nx.spring_layout(G)

def nx_draw_network_nodes_labels(G, pos, _node_size=100, _font_size=10):
    # 绿色 节点
    nx.draw_networkx_nodes(G, pos, node_color='g', node_size = _node_size)
    # 红色 标签
    nx.draw_networkx_labels(G, pos, font_color='r', font_size = _font_size, font_family='sans-serif')

def nx_draw_network_edges(G, pos, _edge_list , weight , _edge_color='b',_style='solid'):
    nx.draw_networkx_edges(G, pos, edgelist = _edge_list, width= weight , edge_color=_edge_color, style=_style)

def plt_show():
    plt.axis('off')
    plt.show()

def plt_networkx_main_first(G):
    pos = nx_get_pos(G)
    nx_draw_network_nodes_labels(G, pos)
    coauthor_edge = get_edgelist_by_weight_equal(G, 2.5)
    nx_draw_network_edges(G, pos, coauthor_edge, weight=2.5, _edge_color='b')
    save_fig(plt, "first_stage_networkx_fig")
    plt_show()

def plt_networkx_main_second(G):
    pos = nx_get_pos(G)
    nx_draw_network_nodes_labels(G, pos)
    coauthor_edge = get_edgelist_by_weight_equal(G, 2.5)
    second_edge = get_edgelist_by_weight_equal(G, 2)
    nx_draw_network_edges(G, pos, coauthor_edge, weight=2.5, _edge_color='b')
    nx_draw_network_edges(G, pos, second_edge, weight=2, _edge_color='r', _style='dashed')
    save_fig(plt, "second_stage_networkx_fig")
    plt_show()

def plt_networkx_main_third(G):
    pos = nx_get_pos(G)
    nx_draw_network_nodes_labels(G, pos)
    # 按照 边的 权重值 取  边
    coauthor_edge = get_edgelist_by_weight_equal(G, 2.5)
    second_edge = get_edgelist_by_weight_equal(G, 2)
    third_edge = get_edgelist_by_weight_equal(G, 1.5)

    # 按照 边的 权值 画边
    # 蓝色 第一种边
    nx_draw_network_edges(G, pos, coauthor_edge, weight=2.5, _edge_color='b')
    # 黄色 第二种边
    nx_draw_network_edges(G, pos, second_edge, weight=2, _edge_color='r', _style='dashed' )
    # 绿色 第三种边
    nx_draw_network_edges(G, pos, third_edge, weight=1.5, _edge_color='k',_style='dotted')

    plt_show()

def save_fig(plt, filename):
    plt.savefig("../matplotlib_fig_dir/" + filename)